package hive;

import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF;
import org.apache.hadoop.hive.serde2.objectinspector.*;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.WritableConstantIntObjectInspector;
import org.apache.hadoop.io.IntWritable;

import java.util.ArrayList;
import java.util.List;

/**
 * Created with hadoop-test.
 * 作者: 周艳钢
 * Date: 2016/11/10 0010
 * Time: 下午 4:51
 * 功能：
 */
@Description(name="forx",
        value = "_FUNC_(start, end, inc) - display int from start to end by inc perio.",
        extended = "Example: \n"
                + " > SELECT _FUNC_(1, 10, 2) FROM src;\n"
                + " > SELECT _FUNC_(1, 10) FROM src;")
public class GenericUDTFFor extends GenericUDTF{

    private IntWritable start;
    private IntWritable end;
    private IntWritable inc;

    private Object[] forwardObj = null;

    @Override
    public StructObjectInspector initialize(StructObjectInspector argOIs) throws UDFArgumentException {
        List<? extends StructField> inputFields = argOIs.getAllStructFieldRefs();
        ObjectInspector[] udtfInputOIs = new ObjectInspector[inputFields.size()];
        for (int i = 0; i < inputFields.size(); i++) {
            udtfInputOIs[i] = inputFields.get(i).getFieldObjectInspector();
        }

        start = ((WritableConstantIntObjectInspector) udtfInputOIs[0]).getWritableConstantValue();
        end = ((WritableConstantIntObjectInspector) udtfInputOIs[1]).getWritableConstantValue();

        if (udtfInputOIs.length == 3) {
            inc = ((WritableConstantIntObjectInspector) udtfInputOIs[2]).getWritableConstantValue();
        } else {
            inc = new IntWritable(1);
        }

        forwardObj = new Object[1];

        ArrayList<String> fieldNames = new ArrayList<String>();
        ArrayList<ObjectInspector> fieldOIs = new ArrayList<ObjectInspector>();

        fieldNames.add("col0");
        fieldOIs.add(
                PrimitiveObjectInspectorFactory.getPrimitiveJavaObjectInspector(PrimitiveObjectInspector.PrimitiveCategory.INT)
        );

        return ObjectInspectorFactory.getStandardStructObjectInspector(fieldNames,
                fieldOIs);
    }

    @Override
    public void process(Object[] args) throws HiveException {

        for(int i = start.get(); i < end.get(); i += inc.get()) {
            forwardObj[0] = new Integer(i);
            forward(forwardObj);
        }
    }

    @Override
    public void close() throws HiveException {

    }
}
